#include "pcbDetect.h"

// 类初始化：构造函数
PCBDetector::PCBDetector(const std::string& model_path, float confidence_threshold, float nms_threshold)
    : confidence_threshold_(confidence_threshold),
      nms_threshold_(nms_threshold),
      env_(nullptr),
      session_(nullptr),
      allocator_(nullptr),
      input_name_(nullptr),
      num_output_nodes_(0) {
    
    // 初始化类别映射
    class_mapping_ = {
        {0, "missing_hole"},   // 漏孔
        {1, "mouse_bite"},     // 鼠咬
        {2, "open_circuit"},   // 开路
        {3, "short"},          // 短路
        {4, "spur"},           // 毛刺
        {5, "spurious_copper"},// 杂铜
        {6, "pinhole"},        // 针孔
        {7, "scratch"}         // 划痕
    };
    
    // 初始化颜色映射
    colors_ = {
        cv::Scalar(0, 255, 0), cv::Scalar(255, 0, 0), cv::Scalar(0, 0, 255),
        cv::Scalar(255, 255, 0), cv::Scalar(255, 0, 255), cv::Scalar(0, 255, 255),
        cv::Scalar(128, 128, 128), cv::Scalar(255, 128, 0)
    };
    
    // 初始化ONNX Runtime
    initOrt(model_path);
}

// 析构函数
PCBDetector::~PCBDetector() {
    // 释放资源
    cleanup();
}

// 从文件路径检测
std::vector<Detection> PCBDetector::detectFromPath(const std::string& image_path, cv::Mat& result_image) {
    // 加载图像
    cv::Mat image = cv::imread(image_path);
    if (image.empty()) {
        std::cerr << "错误: 无法读取图像 " << image_path << std::endl;
        result_image = cv::Mat();
        return {};
    }
    
    return detect(image, result_image);
}

// 从OpenCV Mat检测
std::vector<Detection> PCBDetector::detect(const cv::Mat& input_image, cv::Mat& result_image) {
    if (input_image.empty()) {
        std::cerr << "错误: 输入图像为空" << std::endl;
        result_image = cv::Mat();
        return {};
    }
    
    int orig_width = input_image.cols;
    int orig_height = input_image.rows;
    
    // 预处理图像
    cv::Mat resized, normalized;
    cv::resize(input_image, resized, cv::Size(640, 640));
    resized.convertTo(normalized, CV_32F, 1.0/255.0);
    
    std::vector<cv::Mat> channels(3);
    cv::split(normalized, channels);
    
    std::vector<float> input_data;
    input_data.reserve(3 * 640 * 640);
    
    for (int c = 0; c < 3; ++c) {
        for (int h = 0; h < 640; ++h) {
            for (int w = 0; w < 640; ++w) {
                input_data.push_back(channels[c].at<float>(h, w));
            }
        }
    }
    
    // 创建输入张量
    OrtMemoryInfo* memory_info;
    if (!checkStatus(ort_api_->CreateCpuMemoryInfo(OrtArenaAllocator, OrtMemTypeDefault, &memory_info))) {
        result_image = cv::Mat();
        return {};
    }
    
    const int64_t input_shape[] = {1, 3, 640, 640};
    const size_t input_shape_len = sizeof(input_shape) / sizeof(input_shape[0]);
    
    OrtValue* input_tensor = nullptr;
    if (!checkStatus(ort_api_->CreateTensorWithDataAsOrtValue(memory_info, input_data.data(), 
                                                           input_data.size() * sizeof(float),
                                                           input_shape, input_shape_len, 
                                                           ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, 
                                                           &input_tensor))) {
        ort_api_->ReleaseMemoryInfo(memory_info);
        result_image = cv::Mat();
        return {};
    }
    
    // 运行推理
    OrtValue* output_tensors[8] = {nullptr};
    const char* input_names[] = {input_name_};
    const char* output_names_array[8];
    for (size_t i = 0; i < num_output_nodes_ && i < 8; i++) {
        output_names_array[i] = output_names_[i];
    }
    
    OrtRunOptions* run_options;
    if (!checkStatus(ort_api_->CreateRunOptions(&run_options))) {
        ort_api_->ReleaseMemoryInfo(memory_info);
        ort_api_->ReleaseValue(input_tensor);
        result_image = cv::Mat();
        return {};
    }
    
    if (!checkStatus(ort_api_->Run(session_, run_options, input_names, 
                                 (const OrtValue* const*)&input_tensor, 1,
                                 output_names_array, num_output_nodes_, output_tensors))) {
        ort_api_->ReleaseRunOptions(run_options);
        ort_api_->ReleaseMemoryInfo(memory_info);
        ort_api_->ReleaseValue(input_tensor);
        result_image = cv::Mat();
        return {};
    }
    
    ort_api_->ReleaseRunOptions(run_options);
    
    // 获取输出数据
    float* boxes_data = nullptr;
    float* labels_data = nullptr; 
    float* scores_data = nullptr;
    
    if (!checkStatus(ort_api_->GetTensorMutableData(output_tensors[0], (void**)&boxes_data))) {
        for (size_t i = 0; i < num_output_nodes_; i++) {
            ort_api_->ReleaseValue(output_tensors[i]);
        }
        ort_api_->ReleaseMemoryInfo(memory_info);
        ort_api_->ReleaseValue(input_tensor);
        result_image = cv::Mat();
        return {};
    }
    
    if (!checkStatus(ort_api_->GetTensorMutableData(output_tensors[1], (void**)&labels_data))) {
        for (size_t i = 0; i < num_output_nodes_; i++) {
            ort_api_->ReleaseValue(output_tensors[i]);
        }
        ort_api_->ReleaseMemoryInfo(memory_info);
        ort_api_->ReleaseValue(input_tensor);
        result_image = cv::Mat();
        return {};
    }
    
    if (!checkStatus(ort_api_->GetTensorMutableData(output_tensors[2], (void**)&scores_data))) {
        for (size_t i = 0; i < num_output_nodes_; i++) {
            ort_api_->ReleaseValue(output_tensors[i]);
        }
        ort_api_->ReleaseMemoryInfo(memory_info);
        ort_api_->ReleaseValue(input_tensor);
        result_image = cv::Mat();
        return {};
    }
    
    // 后处理
    std::vector<Detection> detections = postprocessResults(boxes_data, labels_data, scores_data, 8, orig_width, orig_height);
    
    // 绘制结果
    result_image = drawDetections(input_image, detections);
    
    // 清理资源
    for (size_t i = 0; i < num_output_nodes_; i++) {
        ort_api_->ReleaseValue(output_tensors[i]);
    }
    ort_api_->ReleaseMemoryInfo(memory_info);
    ort_api_->ReleaseValue(input_tensor);
    
    return detections;
}

// 初始化ONNX Runtime
void PCBDetector::initOrt(const std::string& model_path) {
    // 获取OrtApi
    ort_api_ = OrtGetApiBase()->GetApi(ORT_API_VERSION);
    
    // 创建环境
    if (!checkStatus(ort_api_->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "PCBDetection", &env_))) {
        throw std::runtime_error("Failed to create ORT environment");
    }
    
    // 创建会话选项
    OrtSessionOptions* session_options;
    if (!checkStatus(ort_api_->CreateSessionOptions(&session_options))) {
        ort_api_->ReleaseEnv(env_);
        throw std::runtime_error("Failed to create session options");
    }
    
    // 创建会话
    if (!checkStatus(ort_api_->CreateSession(env_, model_path.c_str(), session_options, &session_))) {
        ort_api_->ReleaseSessionOptions(session_options);
        ort_api_->ReleaseEnv(env_);
        throw std::runtime_error("Failed to create session");
    }
    ort_api_->ReleaseSessionOptions(session_options);
    
    // 获取输入输出信息
    size_t num_input_nodes;
    if (!checkStatus(ort_api_->SessionGetInputCount(session_, &num_input_nodes))) {
        cleanup();
        throw std::runtime_error("Failed to get input count");
    }
    
    if (!checkStatus(ort_api_->SessionGetOutputCount(session_, &num_output_nodes_))) {
        cleanup();
        throw std::runtime_error("Failed to get output count");
    }
    
    // 获取输入名称
    if (!checkStatus(ort_api_->GetAllocatorWithDefaultOptions(&allocator_))) {
        cleanup();
        throw std::runtime_error("Failed to get allocator");
    }
    
    if (!checkStatus(ort_api_->SessionGetInputName(session_, 0, allocator_, &input_name_))) {
        allocator_->Free(allocator_, input_name_);
        cleanup();
        throw std::runtime_error("Failed to get input name");
    }
    
    // 获取输出名称
    for (size_t i = 0; i < num_output_nodes_; i++) {
        char* output_name;
        if (!checkStatus(ort_api_->SessionGetOutputName(session_, i, allocator_, &output_name))) {
            for (char* name : output_names_) {
                allocator_->Free(allocator_, name);
            }
            allocator_->Free(allocator_, input_name_);
            cleanup();
            throw std::runtime_error("Failed to get output name");
        }
        output_names_.push_back(output_name);
    }
    
    std::cout << "ONNX模型加载成功，输入名称: " << input_name_ << ", 输出节点数: " << num_output_nodes_ << std::endl;
}

// 后处理检测结果
std::vector<Detection> PCBDetector::postprocessResults(float* boxes_data, float* labels_data, float* scores_data, 
                                                     int num_detections, int orig_width, int orig_height) {
    std::vector<Detection> detections;
    
    // 计算缩放比例
    float scale_x = (float)orig_width / 640.0f;
    float scale_y = (float)orig_height / 640.0f;
    
    for (int i = 0; i < num_detections; i++) {
        float score = scores_data[i];
        int class_id = static_cast<int>(labels_data[i]);
        
        auto it = class_mapping_.find(class_id);
        std::string class_name = (it != class_mapping_.end()) ? it->second : "unknown_class_" + std::to_string(class_id);
        
        if (score > confidence_threshold_) {
            // 原始坐标（相对于640x640的）
            float raw_x1 = boxes_data[i * 4 + 0];
            float raw_y1 = boxes_data[i * 4 + 1];
            float raw_x2 = boxes_data[i * 4 + 2];
            float raw_y2 = boxes_data[i * 4 + 3];
            
            // 缩放回原始图像尺寸
            float x1 = raw_x1 * scale_x;
            float y1 = raw_y1 * scale_y;
            float x2 = raw_x2 * scale_x;
            float y2 = raw_y2 * scale_y;
            
            // 确保坐标在有效范围内
            x1 = std::max(0.0f, std::min((float)orig_width, x1));
            y1 = std::max(0.0f, std::min((float)orig_height, y1));
            x2 = std::max(0.0f, std::min((float)orig_width, x2));
            y2 = std::max(0.0f, std::min((float)orig_height, y2));
            
            // 确保 x2 > x1, y2 > y1
            if (x2 <= x1 || y2 <= y1) {
                continue;
            }
            
            cv::Rect bbox(static_cast<int>(x1), static_cast<int>(y1), 
                         static_cast<int>(x2 - x1), static_cast<int>(y2 - y1));
            
            Detection det;
            det.bbox = bbox;
            det.confidence = score;
            det.class_id = class_id;
            det.class_name = class_name;
            
            detections.push_back(det);
        }
    }
    
    return detections;
}

// 绘制检测结果
cv::Mat PCBDetector::drawDetections(const cv::Mat& image, const std::vector<Detection>& detections) {
    cv::Mat result = image.clone();
    
    for (const auto& det : detections) {
        cv::Scalar color = colors_[det.class_id % colors_.size()];
        
        // 绘制边界框
        cv::rectangle(result, det.bbox, color, 8);
        
        // 绘制标签
        std::string label = det.class_name + ": " + std::to_string(det.confidence).substr(0, 5);
        
        // 调整字体大小以适应图像尺寸
        double font_scale = std::max(1.0, (double)image.cols / 1000.0);
        int thickness = std::max(2, (int)(font_scale * 2));
        
        int baseline;
        cv::Size label_size = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, font_scale, thickness, &baseline);
        
        cv::Point label_pos(det.bbox.x, det.bbox.y - 20);
        if (label_pos.y < label_size.height) {
            label_pos.y = det.bbox.y + det.bbox.height + label_size.height + 20;
        }
        
        // 绘制标签背景
        cv::rectangle(result, 
                     cv::Point(label_pos.x, label_pos.y - label_size.height - baseline),
                     cv::Point(label_pos.x + label_size.width, label_pos.y + baseline),
                     color, cv::FILLED);
        
        // 绘制标签文字
        cv::putText(result, label, label_pos, cv::FONT_HERSHEY_SIMPLEX, font_scale, cv::Scalar(0, 0, 0), thickness);
        
        // 在边界框中心绘制类别ID
        cv::Point center(det.bbox.x + det.bbox.width/2, det.bbox.y + det.bbox.height/2);
        cv::putText(result, std::to_string(det.class_id), center, cv::FONT_HERSHEY_SIMPLEX, font_scale, color, thickness);
    }
    
    return result;
}

// 清理资源
void PCBDetector::cleanup() {
    if (allocator_) {
        for (char* name : output_names_) {
            allocator_->Free(allocator_, name);
        }
        if (input_name_) {
            allocator_->Free(allocator_, input_name_);
        }
        allocator_ = nullptr;
    }
    
    if (session_) {
        ort_api_->ReleaseSession(session_);
        session_ = nullptr;
    }
    
    if (env_) {
        ort_api_->ReleaseEnv(env_);
        env_ = nullptr;
    }
    
    output_names_.clear();
    input_name_ = nullptr;
    num_output_nodes_ = 0;
}

// 检查状态
bool PCBDetector::checkStatus(OrtStatus* status) {
    if (status != nullptr) {
        const char* msg = ort_api_->GetErrorMessage(status);
        std::cerr << "ONNX Runtime 错误: " << msg << std::endl;
        ort_api_->ReleaseStatus(status);
        return false;
    }
    return true;
}